from langchain_community.chat_message_histories import ChatMessageHistory
from langchain_core.chat_history import BaseChatMessageHistory
from langchain_core.runnables.history import RunnableWithMessageHistory

from langchain_core.messages import HumanMessage
from langchain_core.runnables import RunnableParallel
from langchain_zhipu import ChatZhipuAI
from operator import itemgetter

model = ChatZhipuAI(api_key="d3708ee404327e207b2f003775e06908.X3dgRCxbkyDfEIbh"
                    , model="glm-4")

chain = RunnableParallel({"output_message": model})

store = {}


def get_session_history(session_id: str) -> BaseChatMessageHistory:
    if session_id not in store:
        store[session_id] = ChatMessageHistory()
    return store[session_id]


with_message_history = RunnableWithMessageHistory(
    chain,
    get_session_history,
    output_messages_key="output_message",
)

with_message_history.invoke(
    [HumanMessage(content="什么是langchain？")],
    config={"configurable": {"session_id": "baz"}},
)
print(with_message_history)

answer = RunnableWithMessageHistory(
    itemgetter("input_messages") | model,
    get_session_history,
    input_messages_key="input_messages",
)
cc = answer.invoke(
    [HumanMessage(content="精简下内容")],
    config={"configurable": {"session_id": "baz"}},
)

print(cc)
